// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/python/runner/PyConnectors.h>
#include <pollux/connectors/hive/hive_connector.h>
#include <pollux/connectors/tpch/tpch_connector.h>

namespace kumo::pollux::py {
namespace {

std::unordered_set<std::string>& connectorRegistry() {
  static std::unordered_set<std::string> registry;
  return registry;
}

template <typename TConnectorFactory>
void registerConnector(
    const std::string& connectorId,
    melon::F14FastMap<std::string, std::string> configs) {
  connector::registerConnectorFactory(
      std::make_shared<TConnectorFactory>(connectorId.data()));

  const auto configBase =
      std::make_shared<pollux::config::ConfigBase>(std::move(configs));
  auto connector =
      connector::getConnectorFactory(connectorId)
          ->newConnector(
              connectorId, configBase, melon::getGlobalCPUExecutor().get());
  connector::registerConnector(connector);
  connectorRegistry().insert(connectorId);
}

} // namespace

void registerHive(
    const std::string& connectorId,
    melon::F14FastMap<std::string, std::string> configs) {
  registerConnector<connector::hive::HiveConnectorFactory>(
      connectorId, std::move(configs));
}

void registerTpch(
    const std::string& connectorId,
    melon::F14FastMap<std::string, std::string> configs) {
  registerConnector<connector::tpch::TpchConnectorFactory>(
      connectorId, std::move(configs));
}

// Is it ok to unregister connectors that were not registered.
void unregister(const std::string& connectorId) {
  if (!kumo::pollux::connector::unregisterConnector(connectorId) ||
      !kumo::pollux::connector::unregisterConnectorFactory(connectorId)) {
    throw std::runtime_error(
        fmt::format("Unable to unregister connector '{}'", connectorId));
  }
  connectorRegistry().erase(connectorId);
}

void unregisterAll() {
  while (!connectorRegistry().empty()) {
    unregister(*connectorRegistry().begin());
  }
}

} // namespace kumo::pollux::py
